#
# fusion.db.sqlite_db
#
# Copyright 2003-2007 Helsinki Institute for Information Technology
# and the authors.
#
# Authors: Ken Rimey <rimey@hiit.fi>
#          Kenneth Oksanen <cessu@iki.fi>
#

# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

# This file derives from db.py in PDIS (http://pdis.hiit.fi/).

# We require APSW:
#
#   http://initd.org/pub/software/pysqlite/apsw/3.3.13-r1/apsw.html#Building
#
# Migrating to the pysqlite included in Python 2.5 would simplify
# installation, but it has subtle differences, so this is not
# something to rush into.

# XXX Support hierarchical groups (includes)?  One idea is to replace
# the creator link with a list of links to all of the items in the
# authorizing certificate chain starting with the rights item and
# ending with the member item.

from binascii import b2a_hex as bytes_to_hex
from binascii import a2b_hex as hex_to_bytes

import os
import hashlib
import logging

from apsw import Connection

from fusion.lib.sexp import encode, decode
from fusion.core.identity import generate_identity, decode_identity
from fusion.core.item import item_from_row, row_from_item, \
    create_root_item, create_kv_item, create_kv0_item, create_kv1_item
from fusion.db.util import *
from fusion.db.exceptions import *

default_start_collection = hex_to_bytes('0c2effbf3ab4f8b6fe2887ceb35ea6689d421395')

INFO_TABLE = """
create table if not exists info(
    k primary key,
    v);"""

COLLECTION_TABLE = """
create table if not exists collections(
    num integer primary key autoincrement,
    id unique);"""

ITEM_TABLE = """
create table if not exists items(
    num integer primary key autoincrement,
    id unique,
    type,
    collection,
    key,
    value,
    size integer,
    sequence integer,
    salt,
    creator,
    signature);

create index if not exists items_by_num
on items(collection, num);

create index if not exists items_by_key
on items(collection, key, num);"""

class SqliteDB(object):
    '''
    Database Instance
    '''
    def __init__(self, path=':memory:'):
        path = os.path.normpath(os.path.expanduser(path))
        dir = os.path.dirname(path)
        if dir and not os.path.exists(dir):
            os.makedirs(dir)

        self.connection = Connection(path)
        self.connection.setbusytimeout(2 * 60 * 1000)

        self.nesting = 0
        self.cursor = self.connection.cursor()

        # This pragma has no effect if I execute it in a transaction.
        self.execute("pragma auto_vacuum=1;")

        # Default is FULL, which is supposedly slow.  However, the
        # following seems to cause problems where the database is left
        # locked.
        #self.execute("pragma synchronous = NORMAL;")

        self.operate(INFO_TABLE + COLLECTION_TABLE + ITEM_TABLE)
        self.init_start_collection()
        self.init_default_identity()

    def close(self):
        self.cursor.close()
        self.connection.close()

    #
    # Database info table
    #

    def get_local_info(self, key):
        c = self.execute("select v from info where k=?;", (key,))
        try:
            value = get_single_value(c)
        except LookupError:
            return None
        return str(value) if isinstance(value, buffer) else value

    def set_local_info(self, key, value):
        if isinstance(value, str):
            value = buffer(value)
        self.operate("replace into info values (?, ?);", (key, value))

    def get_default_identity(self):
        value = self.get_local_info('default_identity')
        return value and decode_identity(value)

    def set_default_identity(self, identity):
        self.set_local_info('default_identity', identity.encode_private())

    def init_default_identity(self, *args, **options):
        def go():
            identity = self.get_default_identity()
            if identity is None:
                identity = generate_identity(*args, **options)
                self.set_default_identity(identity)

        self.exclusive(go)

    def get_start_collection(self):
        return self.get_local_info('start_collection')

    def set_start_collection(self, collection):
        self.set_local_info('start_collection', collection)

    def init_start_collection(self):
        def go():
            c = self.get_start_collection()
            if c is None:
                c = default_start_collection
                self.set_start_collection(c)

            self.import_collection(c)

        self.exclusive(go)

    #
    # Collections
    #

    def list_collections(self):
        c = self.execute("select id from collections;")
        return map(str, get_values(c))

    def collection_exists(self, collection):
        # XXX Should this check that the root item exists?
        id = buffer(collection)
        c = self.execute("select count() from collections"
                         " where id=?;", (id,))
        return bool(get_single_value(c))

    def get_owner(self, collection):
        root = self.get_item(collection)
        if root is None or not root.is_root():
            return None

        return root.get_identity()

    def create_collection(self, my_identity=None):
        my_identity = my_identity or self.get_default_identity()
        root = create_root_item(my_identity)

        def go():
            self.import_collection(root.id)
            num = self.put_item(root)
            assert num
            return root.id

        return self.exclusive(go)

    def import_collection(self, collection):
        assert isinstance(collection, str) and len(collection) == 20
        id = buffer(collection)
        try:
            self.operate("insert into collections(id)"
                         " values (?);", (id,))
        except ConstraintError:
            pass

    def remove_collection(self, collection, force=False):
        def go():
            if not force and self._get_dependent_collections(collection):
                raise ValueError('Cannot remove collection'
                                 ' because another collection requires it.')
            self._remove_collection(collection)

        self.exclusive(go)

    def _remove_collection(self, collection):
        id = buffer(collection)
        self.operate("delete from collections where id=?;"
                     "delete from items where collection=?;",
                     (id, id))

    def _get_dependent_collections(self, collection):
        def gen():
            q = self.query_all_collections(['rights', collection])
            try:
                for c, k, v in q:
                    if v == '1' and c != collection:
                        yield c
            finally:
                q.close()

        return set(gen())

    #
    # Application data access
    #

    def get(self, collection, key):
        collection = buffer(collection)
        key = buffer(encode(key))
        c = self.execute("select value from items"
                         " where collection=? and key=?"
                         " order by num desc",
                         (collection, key))
        try:
            return str(get_single_value(c))
        except LookupError:
            if self.collection_exists(collection):
                return None
            else:
                raise NoSuchCollection

    def get_size(self, collection, key):
        collection = buffer(collection)
        key = buffer(encode(key))
        c = self.execute("select size from items"
                         " where collection=? and key=?"
                         " order by num desc",
                         (collection, key))
        try:
            return get_single_value(c)
        except LookupError:
            if self.collection_exists(collection):
                return None
            else:
                raise NoSuchCollection

    def put(self, collection, key, value, my_identity=None):
        assert isinstance(value, str)
        my_identity = my_identity or self.get_default_identity()
        self._put(collection, key, value, my_identity, op='put')

    def post(self, collection, key, value, my_identity=None):
        assert isinstance(value, str)
        my_identity = my_identity or self.get_default_identity()
        key.append(my_identity.get_fingerprint())
        self._put(collection, key, value, my_identity, op='post')
        return key

    def _put(self, collection, key, value, my_identity, op):
        try:
            return self.__put(collection, key, value, my_identity, op)
        except Unauthorized:
            required_rights = [op] + key
            for rights in self.list_rights(collection, collection, ''):
                if is_prefix(rights, required_rights):
                    self.add_member(collection, '', my_identity, my_identity)
                    return self.__put(collection, key, value, my_identity, op)

            if self.collection_exists(collection):
                raise
            else:
                raise NoSuchCollection

    def __put(self, collection, key, value, my_identity, op):
        def gather_info():
            existing = self.get_item_by_key(collection, key)
            signing = self._find_signing_item(
                collection, [op] + key, my_identity)
            return existing, signing

        # We use two transactions to avoid transacting the signature
        # computation.  If a newer item is inserted into the database
        # in the interim, the put_item call will have no effect, but
        # this case is indistinguishable from the case where the item
        # that we insert is overwritten immediately afterwards.

        existing_item, signing_item = self.transact(gather_info)

        if existing_item is None or existing_item.value != value:
            if existing_item is None:
                sequence = 1
            else:
                sequence = existing_item.seq + 1

            if signing_item is not None:
                if signing_item.is_root():
                    item = create_kv_item(collection, key, value, sequence,
                                          my_identity)
                else:
                    item = create_kv1_item(collection, key, value, sequence,
                                           signing_item.id, my_identity)
            else:
                if key[0] == 'member' and key[1] == '':
                    item = create_kv0_item(collection, key, value, sequence,
                                           my_identity)
                else:
                    raise Unauthorized

            self.put_item(item)

    def _find_signing_item(self, collection, rights, identity):
        root = self.get_item(collection)
        if root is None or not root.is_root():
            return None

        if root.value == identity.encode_public():
            return root

        # XXX If we need post rights, should we also check for
        # nominations with put rights?
        for group, role in self.list_authorized_groups(collection, rights):
            item = self._get_nomination(group, role, identity)
            if item is not None and item.value:
                return item

        return None

    #
    # Groups
    #

    def add_member(self, group, role, member_identity, my_identity=None):
        my_identity = my_identity or self.get_default_identity()
        key = ['member', role, member_identity.get_fingerprint()]
        value = member_identity.encode_public()
        self.put(group, key, value, my_identity)

    def remove_member(self, group, role, member_identity, my_identity=None):
        my_identity = my_identity or self.get_default_identity()
        key = ['member', role, member_identity.get_fingerprint()]
        self.put(group, key, '', my_identity)

    def is_member(self, group, role, identity):
        key = ['member', role, identity.get_fingerprint()]
        return bool(self.get(group, key))

    def _get_nomination(self, group, role, identity):
        key = ['member', role, identity.get_fingerprint()]
        return self.get_item_by_key(group, key)

    def list_members(self, group, role=None):
        def gen():
            prefix = ['member']
            if role is not None:
                prefix.append(role)

            q = self.query(group, prefix)
            try:
                for key, value in q:
                    if value:
                        try:
                            identity = decode_identity(value)
                        except ValueError:
                            continue

                        if role is not None:
                            yield identity
                        else:
                            yield identity, key[0]
            finally:
                q.close()

        return list(gen())

    #
    # Rights
    #

    def add_rights(self, collection, group, role, rights, my_identity=None):
        key = ['rights', group, role] + rights
        self.put(collection, key, '1', my_identity)

    def remove_rights(self, collection, group, role, rights, my_identity=None,
                      retroactive=False):
        key = ['rights', group, role] + rights
        value = '' if retroactive else '0'
        self.put(collection, key, value, my_identity)

    def list_rights(self, collection, group, role):
        key_prefix = ['rights', group, role]
        q = self.query(collection, key_prefix)
        try:
            return [k for k, v in q if v == '1']
        finally:
            q.close()

    def has_rights(self, collection, rights, identity=None):
        # XXX Also answer affirmatively if the collection grants these
        # rights to its guests?
        identity = identity or self.get_default_identity()
        item = self.transact(self._find_signing_item,
                             collection, rights, identity)
        return item is not None

    def list_authorized_groups(self, collection, required_rights):
        def gen():
            q = self.query(collection, ['rights'])
            try:
                for key, value in q:
                    if value == '1':
                        rights = key[2:]
                        if is_prefix(rights, required_rights):
                            group, role = key[:2]
                            yield len(rights), group, role
            finally:
                q.close()

        # We order the list to favor roles granted narrower rights.
        candidates = sorted(gen(), reverse=True)
        return [(group, role) for goodness, group, role in candidates]

    #
    # Queries
    #

    def get_mark(self):
        c = self.execute("select seq from sqlite_sequence"
                         " where name='items';")
        try:
            return get_single_value(c)
        except LookupError:
            return 0

    def keys(self, collection, key_prefix=None, last_mark=0, reverse=False):
        query = "select num, key from items where collection = ?"
        params = [buffer(collection)]

        prefix = ''
        if key_prefix:
            prefix = encode(key_prefix)[:-1]
            query += " and key > ?"
            params += [buffer(prefix)]

        if last_mark:
            query += " and num > ?"
            params += [last_mark]

        query += " order by key asc, num desc;"

        results = []
        last_key = None
        for num, key in self.execute(query, params):
            if key != last_key:
                last_key = key

                key = str(key)
                if key.startswith(prefix):
                    chop = 0 if not key_prefix else len(key_prefix)
                    result_key = decode(key)[chop:]
                    results.append((num, result_key))
                else:
                    break           # Skip all the rest.

        if not results and not self.collection_exists(collection):
            raise NoSuchCollection

        results.sort(reverse=reverse)
        return [key for num, key in results]

    def query(self, collection, key_prefix=None, last_mark=0,
              most_recent_only=True):
        query = "select key, value from items where collection = ?"
        params = [buffer(collection)]

        prefix = ''
        if key_prefix:
            prefix = encode(key_prefix)[:-1]
            query += " and key > ?"
            params += [buffer(prefix)]

        if last_mark:
            query += " and num > ?"
            params += [last_mark]

        if most_recent_only:
            query += " order by key asc, num desc;"
        else:
            query += " order by key asc, num asc;"

        hits = 0
        cursor = self.execute_with_new_cursor(query, params)
        try:
            last_k = None
            for row in cursor:
                k, v = row
                if k != last_k:
                    last_k = k
                elif most_recent_only:
                    continue        # Skip older entries.

                k = str(k)
                if k.startswith(prefix):
                    hits += 1
                    n = 0 if not key_prefix else len(key_prefix)
                    yield decode(k)[n:], str(v)
                else:
                    break           # Skip all the rest.
        finally:
            cursor.close()

        if not hits and not self.collection_exists(collection):
            raise NoSuchCollection

    def query_all_collections_and_keys(self, last_mark=0):
        cursor = self.execute_with_new_cursor(
            "select collection, key from items"
            " where num > ? order by num;",
            (last_mark,))
        try:
            for row in cursor:
                c, k = map(str, row)
                yield c, decode(k)
        finally:
            cursor.close()

    def query_all_collections(self, key_prefix,
                              last_mark=0, current_mark=None):
        query = "select collection, key, value from items"

        prefix = encode(key_prefix)[:-1]
        query += " where key > ?"
        params = [buffer(prefix)]

        if last_mark:
            query += " and num > ?"
            params += [last_mark]

        if current_mark is not None:
            query += " and num <= ?"
            params += [current_mark]

        query += " order by key, num;"

        cursor = self.execute_with_new_cursor(query, params)
        try:
            for row in cursor:
                c, k, v = map(str, row)
                if k.startswith(prefix):
                    yield c, decode(k), v
                else:
                    break           # Skip all the rest.
        finally:
            cursor.close()

    #
    # Numbers
    #

    def query_ids(self, collection, last_mark=0, current_mark=None):
        query = "select num, id, key from items where collection = ?"
        params = [buffer(collection)]

        if last_mark:
            query += " and num > ?"
            params += [last_mark]

        if current_mark is not None:
            query += " and num <= ?"
            params += [current_mark]

        query += " order by num;"

        cursor = self.execute_with_new_cursor(query, params)
        try:
            prefix = encode(['rights'])[:-1]
            for num, id, key in cursor:
                id = str(id)
                key = str(key)

                group = None
                if key.startswith(prefix):
                    item = self.get_item(num)
                    assert item.is_rights()
                    if item.value == '1':
                        group = item.key[1]

                yield num, str(id), group
        finally:
            cursor.close()

    def num_to_id(self, num):
        c = self.execute("select id from items"
                         " where num=?", (num,))
        try:
            return str(get_single_value(c))
        except LookupError:
            return None

    def id_to_num(self, id):
        id = buffer(id)
        c = self.execute("select num from items"
                         " where id=?", (id,))
        try:
            return get_single_value(c)
        except LookupError:
            return None

    #
    # Item access
    #

    def get_item_by_key(self, collection, key):
        collection = buffer(collection)
        key = buffer(encode(key))
        c = self.execute("select * from items"
                         " where collection=? and key=?"
                         " order by num desc",
                         (collection, key))
        try:
            row = get_single_row(c)
        except LookupError:
            return None

        return item_from_row(row)

    def get_item(self, id_or_num):
        if isinstance(id_or_num, str):
            id = buffer(id_or_num)
            c = self.execute("select * from items"
                             " where id=?", (id,))
        elif isinstance(id_or_num, (int, long)):
            num = id_or_num
            c = self.execute("select * from items"
                             " where num=?", (num,))
        else:
            raise ValueError

        try:
            row = get_single_row(c)
        except LookupError:
            return None

        return item_from_row(row)

    def get_signing_item(self, item):
        if item.creator == item.id:
            return item
        else:
            return self.get_item(item.creator)

    def put_item(self, item):
        signing_item = self.get_signing_item(item)

        if not item.verify_signature(signing_item):
            if not self.collection_exists(item.cn):
                raise NoSuchCollection
            raise Unauthorized

        def go(item):
            if not self._verify_authorization(item, signing_item):
                if not self.collection_exists(item.cn):
                    raise NoSuchCollection
                raise Unauthorized

            existing_item = self.get_item_by_key(item.cn, item.key)
            if existing_item is not None:
                if item == existing_item:
                    return existing_item.num # Already have it
                elif item < existing_item:
                    # If this is a rights or non-guest member item, we
                    # sort of actually do want it, even if we already
                    # have a newer version.  Oh well.
                    return None              # Already have a newer one.
                else:
                    # Remove older existing items, unless required for
                    # verification purposes.
                    assert not item.is_root()
                    if not item.is_member() and not item.is_rights():
                        self._delete_by_num(existing_item.num)
                    if item.is_guest_member():
                        self._delete_by_num(existing_item.num)
                    elif item.is_rights() and item.value == '':
                        self._delete_by_key(item.cn, item.key)

            self.operate("insert into items values (?,?,?,?,?,?,?,?,?,?,?);",
                         row_from_item(item))
            num = self.connection.last_insert_rowid()

            # If the database contains an item that grants rights to
            # some group, we require that the group exist.
            if item.is_rights() and item.value == '1':
                group = item.key[1]
                if group != item.cn:
                    self.import_collection(group)

            return num

        return self.exclusive(go, item)

    def _delete_by_num(self, num):
        self.operate("delete from items where num=?", (num,))

    def _delete_by_key(self, collection, key):
        collection = buffer(collection)
        key = buffer(encode(key))
        self.operate("delete from items where collection=? and key=?",
                     (collection, key))

    def verify(self, item):     # Only used in testing.
        signing_item = self.get_signing_item(item)
        return (item.verify_signature(signing_item)
                and self._verify_authorization(item, signing_item))

    def _verify_authorization(self, item, signing_item):
        # This part of the verification must be done within the
        # transaction that inserts the item into the database.

        if item.type == 'root':
            return self.collection_exists(item.cn)

        elif item.type == 'kv':
            return signing_item.is_root()

        elif item.type == 'kv0':
            return self.list_rights(item.cn, item.cn, '')

        elif item.type == 'kv1':
            if not signing_item.is_member():
                return False

            type, role, fingerprint = signing_item.key
            group = signing_item.cn

            # Check that the membership is current.
            if not self.get(group, signing_item.key):
                return False

            # Now check that the member has the required rights.  This
            # incidentally confirms that the collection exists.
            rights_list = self.list_rights(item.cn, group, role)

            required_rights = ['put'] + item.key
            for rights in rights_list:
                if is_prefix(rights, required_rights):
                    return True

            required_rights = ['post'] + item.key
            for rights in rights_list:
                if is_prefix(rights, required_rights):
                    fingerprint = hashlib.sha1(signing_item.value).digest()
                    return item.key[-1] == fingerprint

            return False

        else:
            return False

    #
    # Database helpers
    #

    def transact(self, fun, *args, **keys):
        self.begin()
        return self._finish(fun, args, keys)

    def exclusive(self, fun, *args, **keys):
        self.begin(exclusive=True)
        return self._finish(fun, args, keys)

    def _finish(self, fun, args, keys):
        try:
            result = fun(*args, **keys)
        except:
            self._rollback()
            raise
        else:
            self.commit()
            return result

    def begin(self, exclusive=False):
        if self.nesting == 0:
            if exclusive:
                self.execute("begin exclusive;")
            else:
                self.execute("begin;")
        self.nesting += 1

    def _rollback(self):
        try:
            self.rollback()
        except SQLError:
            logging.warning('Ignoring exception raised by rollback attempt.')

    def rollback(self):
        self.nesting -= 1
        if self.nesting == 0:
            self.execute("rollback;")

    def commit(self):
        self.nesting -= 1
        if self.nesting == 0:
            self.execute("commit;")

    def execute_with_new_cursor(self, statements, bindings=()):
        return self.execute(statements, bindings, new_cursor=True)

    def execute(self, statements, bindings=(), new_cursor=False):
        if new_cursor:
            cursor = self.connection.cursor()
        else:
            cursor = self.cursor

        while True:
            try:
                return cursor.execute(statements, bindings)
            except SchemaChangeError:
                continue

    def operate(self, statements, bindings=()):
        return self.exclusive(self.execute, statements, bindings)

def is_prefix(k1, k2):
    return k1 == k2[:len(k1)]
